import functools
import json
import sys

import cv2
import os
import numpy as np
import math

save_txt_path = 'save.json'

def get_json(img_list):
    rect = []
    rects = []
    data = []
    def cmpr(a,b):
        if (abs(a[1] - b[1]) < 10 and a[0] > b[0]) or a[1] > b[1]:
            return 1
        elif abs(a[1] - b[1]) < 10 and abs(a[0] - b[0]) < 20:
            return 0
        else:
            return -1

    def cmp(a,b):
        if (abs(a[1] - b[1]) < 10 and a[0] > b[0]) or a[1] > b[1]:
            return 1
        elif abs(a[1] - b[1]) < 10 and abs(a[0] - b[0]) < 20:
            return 0
        else:
            return -1

    for img in img_list:
        if img is None:
            print("无法读取图片")
            exit(-1)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        blur = cv2.GaussianBlur(gray, (5, 5), 1.5)
        _, binary = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        thresh, binary = cv2.threshold(blur, int(_ * 0.95), 255, cv2.THRESH_BINARY)

        horizon_k = int(math.sqrt(binary.shape[1]) * 1.2)  # w
        # hors_k = int(binary.shape[1]/ 16)  # w
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (horizon_k, 1))  # 设置内核形状
        horizon = ~cv2.dilate(binary, kernel, iterations=1)  # 膨胀
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (int(horizon_k / 0.9), 1))
        horizon = cv2.dilate(horizon, kernel, iterations=1)

        cnts = cv2.findContours(horizon, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        border_y, border_x = horizon.shape
        # 去除邻近上边界和下边界检测的轮廓
        for cnt in cnts[0]:
            x, y, w, h = cv2.boundingRect(cnt)
            if y < 4 or y > border_y - 8:
                continue
            rects.append([x, y, w, h])
        # 排序
        rects = sorted(rects, key=functools.cmp_to_key(cmpr))

        pre = None
        idx_lst = []
        # 标记不相关的轮廓
        for idx, cnt in enumerate(rects):
            x, y, w, h = cnt
            if w < 150:
                continue
            if pre is None:
                pre = [x, y, w]
            elif 6 < abs(pre[1] - (y + h / 2)) < 70:  # and 10 < abs(pre[0] - x) < pre[2]
                continue
            pre[1] = y + h / 2
            pre[0] = x
            pre[2] = w
            idx_lst.append(idx)

        # 再次筛选
        rects = [rects[x] for x in idx_lst]
        rects = sorted(rects, key=functools.cmp_to_key(cmp))
        # 将检测的水平线扩充成矩形框
        pre_y, pre_h = -1, -1
        for idx, cnt in enumerate(rects):
            x, y, w, h = cnt
            if pre_h == -1:
                pre_y = y
                h = y - 5
                y = 5
                pre_h = h
            else:
                if abs(pre_y - y) < 10:
                    h = pre_h
                    y = max(y - h, 0)
                else:
                    pre_h = abs(y - pre_y) - 10
                    pre_y = y
                    h = pre_h
                    y = pre_y - h
            rects[idx] = [x, y, w, h + 15]
            x = []
            for idx, rect in enumerate(rects):
                x.append({idx: rect})
            data.append(x)
        rects.clear()
    return data

if __name__ == '__main__':
    path = "testdata"
    imgs = []
    file_names = os.listdir(path)
    for file_name in file_names[:10]:
        img_path = os.path.join(path, file_name)
        img = cv2.imread(img_path)
        if img is None:
            print('无法读取图片')
            exit(-1)
        imgs += [img]
    data = get_json(imgs)
    save_js = open(save_txt_path,'w')
    json.dump(data, save_js)
    save_js.close()
